package com.mediocre.eg7publisherconfirms;

import com.mediocre.constant.Constants;
import com.mediocre.util.CommonUtils;
import com.rabbitmq.client.BuiltinExchangeType;
import com.rabbitmq.client.Channel;
import com.rabbitmq.client.ConfirmCallback;
import com.rabbitmq.client.Connection;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.UUID;
import java.util.concurrent.ConcurrentNavigableMap;
import java.util.concurrent.ConcurrentSkipListMap;
import java.util.concurrent.TimeoutException;
import java.util.function.BooleanSupplier;

/**
 * Reliable publishing with publisher confirms.
 *
 * Publisher confirms are a RabbitMQ extension to implement reliable publishing.
 * When publisher confirms are enabled on a channel, messages the client publishes are confirmed asynchronously by the broker,
 * meaning they have been taken care of on the server side.
 * 注意：这里的消息被发送后，得到确认，仅仅是表明该消息到了消息中间件broker，并被broker处理了，并不表明消息已经到了消费者端并被消费者处理完成。
 *
 * 生产者
 */
public class PublisherConfirms {

    private static final int MESSAGE_COUNT = 10;

    public static void main(String[] args) throws IOException, InterruptedException, TimeoutException {
        publishMessagesIndividually();
        publishMessagesInBatch();
        handlePublishConfirmsAsynchronously();
    }

    private static void publishMessagesIndividually() throws IOException, TimeoutException, InterruptedException {
        try (final Connection connection = CommonUtils.createConnection();
             final Channel channel = connection.createChannel()) {
            String queue = UUID.randomUUID().toString();
            channel.queueDeclare(queue, false, false, true, null);
            channel.confirmSelect();

            channel.exchangeDeclare(Constants.EXCHANGE_NAME_DIRECT, BuiltinExchangeType.DIRECT);

            long start = System.nanoTime();
            for (int i = 0; i < MESSAGE_COUNT; i++) {
                String body = String.valueOf(i);
                // channel.basicPublish("", queue, null, body.getBytes(StandardCharsets.UTF_8));
                channel.basicPublish(Constants.EXCHANGE_NAME_DIRECT, Constants.ROUTING_KEY, null, body.getBytes(StandardCharsets.UTF_8));
                channel.waitForConfirmsOrDie(5_000);
            }
            long end = System.nanoTime();
            System.out.format("Published %,d messages individually in %,d ms%n", MESSAGE_COUNT, Duration.ofNanos(end - start).toMillis());
        }
    }

    private static void publishMessagesInBatch() throws IOException, TimeoutException, InterruptedException {
        try (final Connection connection = CommonUtils.createConnection();
             final Channel channel = connection.createChannel()){

            String queue = UUID.randomUUID().toString();
            channel.queueDeclare(queue, false, false, true, null);
            channel.confirmSelect();

            int batchSize = 100;
            int outstandingMessageCount = 0;
            long start = System.nanoTime();

            for (int i = 0; i < MESSAGE_COUNT; i++) {
                String body = String.valueOf(i);
                channel.basicPublish(Constants.EXCHANGE_NAME_DIRECT, Constants.ROUTING_KEY, null, body.getBytes(StandardCharsets.UTF_8));
                outstandingMessageCount++;

                if (outstandingMessageCount == batchSize) {
                    channel.waitForConfirmsOrDie(5_000);
                    outstandingMessageCount = 0;
                }
            }

            if (outstandingMessageCount > 0) {
                channel.waitForConfirmsOrDie(5000);
            }

            long end = System.nanoTime();
            System.out.format("Published %,d messages in batch in %,d ms%n", MESSAGE_COUNT, Duration.ofNanos(end - start).toMillis());
        }
    }

    private static void handlePublishConfirmsAsynchronously() throws IOException, TimeoutException, InterruptedException {
        try (final Connection connection = CommonUtils.createConnection();
             final Channel channel = connection.createChannel()){

            String queue = UUID.randomUUID().toString();
            channel.queueDeclare(queue, false, false, true, null);
            channel.confirmSelect();

            ConcurrentNavigableMap<Long, String> outstandingConfirms = new ConcurrentSkipListMap<>();

            ConfirmCallback cleanOutstandingConfirms = (sequenceNumber, multiple) -> {
                if (multiple) {
                    ConcurrentNavigableMap<Long, String> confirmed = outstandingConfirms.headMap(sequenceNumber, true);
                    confirmed.clear();
                } else {
                    outstandingConfirms.remove(sequenceNumber);
                }
            };

            channel.addConfirmListener(cleanOutstandingConfirms, (sequenceNumber, multiple) -> {
                String body = outstandingConfirms.get(sequenceNumber);
                System.err.format("Message with body %s has been nack-ed. Sequence number: %d, multiple: %b%n",
                        body, sequenceNumber, multiple);
                cleanOutstandingConfirms.handle(sequenceNumber, multiple);
            });

            long start = System.nanoTime();
            for (int i = 0; i < MESSAGE_COUNT; i++) {
                String body = String.valueOf(i);
                outstandingConfirms.put(channel.getNextPublishSeqNo(), body);
                channel.basicPublish(Constants.EXCHANGE_NAME_DIRECT, Constants.ROUTING_KEY, null, body.getBytes(StandardCharsets.UTF_8));
            }

            if (!waitUntil(Duration.ofSeconds(60), () -> outstandingConfirms.isEmpty())) {
                throw new IllegalStateException("All messages could not be confirmed in 60 seconds");
            }

            long end = System.nanoTime();
            System.out.format("Published %,d messages and handled confirms asynchronously in %,d ms%n", MESSAGE_COUNT, Duration.ofNanos(end - start).toMillis());
        }
    }

    private static boolean waitUntil(Duration timeout, BooleanSupplier condition) throws InterruptedException {
        int waited = 0;
        while (!condition.getAsBoolean() && waited < timeout.toMillis()) {
            Thread.sleep(100L);
            waited += 100;
        }
        return condition.getAsBoolean();
    }
}
